#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

用法：
    python3 plot_roll_pitch.py  path/to/file.csv
    python3 plot_roll_pitch.py  path/to/folder
"""
import csv
import argparse
import matplotlib.pyplot as plt
import numpy as np

def plot_signals(csv_path: str, title=None):
    # 读 CSV（第一行即列名）
    with open(csv_path, 'r') as file:
        reader = csv.DictReader(file)
        data = {key: [] for key in reader.fieldnames}
        for row in reader:
            for key, value in row.items():
                data[key].append(float(value))
    # 定义 torque2torquebody 矩阵
    torque2torquebody = np.array([
        [1, -0.5, -0.5],
        [0, np.sqrt(3)/2, -np.sqrt(3)/2],
        [-1, -1, -1]
    ]) * np.sqrt(2) / 2

    # 从 DataFrame 中提取电流数据
    torque_body_fb = np.array([
        data["Torque_body_fb_0"],
        data["Torque_body_fb_1"],
        data["Torque_body_fb_2"]
    ])  # shape: (3, N)

    # 计算 torque_body = torque2torquebody @ currents
    torque_single_motor = np.linalg.inv(torque2torquebody) @ torque_body_fb  # shape: (3, N)

    # 时间轴：行号 × 0.002 s
    time = np.array(data['seconds_from_epoch']) + np.array(data['nanoseconds']) * 1e-9

    # ---------- 图 Torque ----------
    fig1, ax = plt.subplots(3, 1, sharex=True, figsize=(10, 8))
    if title:
        fig1.suptitle(title)
    else:
        fig1.suptitle("Torque-related Signals")

    ax[0].plot(time, data["sensor_data.motor_rpm_speeds_[0]"], label="motor_rpm_speeds_1")
    ax[0].plot(time, data["sensor_data.motor_rpm_speeds_[1]"], label="motor_rpm_speeds_2")
    ax[0].plot(time, data["sensor_data.motor_rpm_speeds_[2]"], label="motor_rpm_speeds_3")
    ax[0].set_ylabel("r/min")
    ax[0].legend(loc='right')
    ax[0].grid(True)

    ax[1].plot(time, data["Torque_0"], label="Torque_1_control")
    ax[1].plot(time, data["Torque_1"], label="Torque_2_control")
    ax[1].plot(time, data["Torque_2"], label="Torque_3_control")
    # ax[1].plot(time, torque_single_motor[0], linestyle="--", label="motor_1_fb")
    # ax[1].plot(time, torque_single_motor[1], linestyle="--", label="motor_2_fb")
    # ax[1].plot(time, torque_single_motor[2], linestyle="--", label="motor_3_fb")
    ax[1].set_ylabel("N*m")
    ax[1].legend(loc='right')
    ax[1].grid(True)

    eps = 1e-8
    scale_1 = np.array(data["Torque_0"]) / (np.array(data["sensor_data.motor_rpm_speeds_[0]"]) + eps)
    scale_2 = np.array(data["Torque_1"]) / (np.array(data["sensor_data.motor_rpm_speeds_[1]"]) + eps)
    scale_3 = np.array(data["Torque_2"]) / (np.array(data["sensor_data.motor_rpm_speeds_[2]"]) + eps)
    ax[2].plot(time, scale_1, label="scale_1")
    ax[2].plot(time, scale_2, label="scale_2")
    ax[2].plot(time, scale_3, label="scale_3")
    ax[2].set_ylabel("Torque/rpm")
    ax[2].legend(loc='right')
    ax[2].grid(True)
    ax[2].set_xlabel("time [s]")

    # ax[3].plot(time, data["body_rpy_0"], marker='x', markersize=4, label="roll")
    # ax[3].plot(time, data["body_rpy_1"], marker='x', markersize=4, label="pitch")
    # ax[3].plot(time, data["rpy_setpoint_0"], linestyle="--", label="setpoint_roll")
    # ax[3].plot(time, data["rpy_setpoint_1"], linestyle="--", label="setpoint_pitch")
    # ax[3].set_ylabel("rad")
    # ax[3].legend(loc='right')
    # ax[3].grid(True)
    #
    # ax[4].plot(time, data["gyro_corrected_0"], marker='x', markersize=4, label="gyro_roll")
    # ax[4].plot(time, data["gyro_corrected_1"], marker='x', markersize=4, label="gyro_pitch")
    # ax[4].set_ylabel("rad")
    # ax[4].legend(loc='right')
    # ax[4].grid(True)

    # 不在此处show，由主程序统一show


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plot roll/pitch related curves from CSV log.")
    parser.add_argument("path", help="Path to CSV file or folder")
    args = parser.parse_args()

    import os
    import matplotlib.pyplot as plt
    try:
        if os.path.isfile(args.path) and args.path.endswith('.csv'):
            plot_signals(args.path, title=os.path.basename(args.path))
            print(f"Plotted signals from {args.path}")
        elif os.path.isdir(args.path):
            csv_files = [f for f in os.listdir(args.path) if f.endswith('.csv')]
            if not csv_files:
                print(f"No CSV files found in folder: {args.path}")
            for csv_file in csv_files:
                csv_path = os.path.join(args.path, csv_file)
                plot_signals(csv_path, title=csv_file)
        else:
            print(f"Invalid path: {args.path}")
        plt.show()
    except KeyboardInterrupt:
        print("\nKeyboardInterrupt received, closing all figures...")
        plt.close('all')